package org.springframework.session.web.socket.handler;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;

import org.springframework.context.ApplicationEvent;
import org.springframework.context.ApplicationEventPublisher;
import org.springframework.session.Session;
import org.springframework.session.web.socket.events.SessionConnectEvent;
import org.springframework.util.Assert;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.WebSocketHandlerDecorator;
import org.springframework.web.socket.handler.WebSocketHandlerDecoratorFactory;

public final class WebSocketConnectHandlerDecoratorFactory implements WebSocketHandlerDecoratorFactory {

	private static final Log logger = LogFactory.getLog(WebSocketConnectHandlerDecoratorFactory.class);

	private final ApplicationEventPublisher eventPublisher;

	public WebSocketConnectHandlerDecoratorFactory(ApplicationEventPublisher eventPublisher) {
		Assert.notNull(eventPublisher, "eventPublisher cannot be null");
		this.eventPublisher = eventPublisher;
	}

	@Override
	public WebSocketHandler decorate(WebSocketHandler handler) {
		return new SessionWebSocketHandler(handler);
	}

	private final class SessionWebSocketHandler extends WebSocketHandlerDecorator {

		SessionWebSocketHandler(WebSocketHandler delegate) {
			super(delegate);
		}

		@Override
		public void afterConnectionEstablished(WebSocketSession wsSession) throws Exception {
			super.afterConnectionEstablished(wsSession);
			publishEvent(new SessionConnectEvent(this, wsSession));
		}

		private void publishEvent(ApplicationEvent event) {
			try {
				WebSocketConnectHandlerDecoratorFactory.this.eventPublisher.publishEvent(event);
			} catch (Throwable ex) {
				logger.error("Error publishing " + event + ".", ex);
			}
		}
	}
}
